package com.ywy.lucene;

import com.ywy.entity.Res;
import com.ywy.utils.DateUtil;
import com.ywy.utils.StringUtil;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.cn.smart.SmartChineseAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.StringField;
import org.apache.lucene.document.TextField;
import org.apache.lucene.index.*;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.*;
import org.apache.lucene.search.highlight.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.StringReader;
import java.nio.file.Paths;
import java.util.Date;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.locks.ReentrantLock;

/**
 * lucene索引
 */
@Component("resIndex")
public class ResIndex {
    private Directory dir = null;

    @Value("${system.lucenePath}")
    private String lucenePath;

    /**
     * 获取IndexWriter实例
     *
     * @return
     * @throws Exception
     */
    public IndexWriter getWriter() throws Exception {
        dir = FSDirectory.open(Paths.get(lucenePath));
        SmartChineseAnalyzer analyzer = new SmartChineseAnalyzer();
        IndexWriterConfig iwc = new IndexWriterConfig(analyzer);
        IndexWriter writer = new IndexWriter(dir, iwc);
        return writer;
    }

    /**
     * 添加文章索引
     *
     * @param res
     */
    public boolean addIndex(Res res) {
        ReentrantLock lock = new ReentrantLock();
        lock.lock();
        try {
            IndexWriter writer = getWriter();
            Document doc = new Document();
            doc.add(new StringField("id", String.valueOf(res.getId()), Field.Store.YES));
            doc.add(new TextField("name", res.getName(), Field.Store.YES));
            doc.add(new StringField("publishDate", DateUtil.formatDate(new Date(), "yyyy-MM-dd"), Field.Store.YES));
            doc.add(new TextField("content", res.getContent(), Field.Store.YES));
            writer.addDocument(doc);
            writer.close();
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        } finally {
            lock.unlock();
        }
        return true;
    }

    /**
     * 更新资源索引
     *
     * @param res
     */
    public boolean updateIndex(Res res) {
        ReentrantLock lock = new ReentrantLock();
        lock.lock();
        try {
            IndexWriter writer = getWriter();
            Document doc = new Document();
            doc.add(new StringField("id", String.valueOf(res.getId()), Field.Store.YES));
            doc.add(new TextField("name", res.getName(), Field.Store.YES));
            doc.add(new StringField("publishDate", DateUtil.formatDate(new Date(), "yyyy-MM-dd"), Field.Store.YES));
            doc.add(new TextField("content", res.getContent(), Field.Store.YES));
            writer.updateDocument(new Term("id", String.valueOf(res.getId())), doc);
            writer.close();
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        } finally {
            lock.unlock();
        }
        return true;
    }

    /**
     * 删除资源索引
     *
     * @param id
     */
    public void deleteIndex(String id) {
        ReentrantLock lock = new ReentrantLock();
        lock.lock();
        try {
            IndexWriter writer = getWriter();
            writer.deleteDocuments(new Term("id", id));
            writer.forceMergeDeletes(); // 强制删除
            writer.commit();
            writer.close();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            lock.unlock();
        }
    }

    /**
     * 分词搜索资源信息，高亮
     * @param q 查询关键字
     * @return
     * @throws Exception
     */
    public List<Res> search(String q) throws Exception {
        dir = FSDirectory.open(Paths.get(lucenePath));
        IndexReader reader = DirectoryReader.open(dir);
        IndexSearcher is = new IndexSearcher(reader);

        BooleanQuery.Builder booleanQuery = new BooleanQuery.Builder();
        SmartChineseAnalyzer analyzer = new SmartChineseAnalyzer();
        QueryParser parserName = new QueryParser("name", analyzer);
        Query queryName = parserName.parse(q);
        booleanQuery.add(queryName, BooleanClause.Occur.SHOULD);

        QueryParser parserContent = new QueryParser("content", analyzer);
        Query queryContent = parserContent.parse(q);
        booleanQuery.add(queryContent, BooleanClause.Occur.SHOULD);

        QueryScorer scorer = new QueryScorer(queryName);
        Fragmenter fragmenter = new SimpleSpanFragmenter(scorer);
        SimpleHTMLFormatter simpleHTMLFormatter = new SimpleHTMLFormatter("<b><label style='color:red;'>", "</label></b>");
        Highlighter highlighter = new Highlighter(simpleHTMLFormatter, scorer);
        highlighter.setTextFragmenter(fragmenter);

        TopDocs hits = is.search(booleanQuery.build(), 100);
        List<Res> resList = new LinkedList<>();
        for (ScoreDoc scoreDoc : hits.scoreDocs) {
            Document doc = is.doc(scoreDoc.doc);
            Res res = new Res();
            res.setId(Integer.parseInt(doc.get(("id"))));
            res.setPublishDateStr(doc.get(("publishDate")));
            String name = doc.get("name");
            String content = StringUtil.stripHtml(doc.get("content"));
            if (name != null) {
                TokenStream tokenStream = analyzer.tokenStream("name", new StringReader(name));
                String hName = highlighter.getBestFragment(tokenStream, name);
                if (StringUtil.isEmpty(hName)) {
                    res.setName(name);
                } else {
                    res.setName(hName);
                }
            }
            if (content != null) {
                TokenStream tokenStream = analyzer.tokenStream("content", new StringReader(content));
                String hContent = highlighter.getBestFragment(tokenStream, content);
                if (StringUtil.isEmpty(hContent)) {
                    if (content.length() <= 200) {
                        res.setContent(content);
                    } else {
                        res.setContent(content.substring(0, 200));
                    }
                } else {
                    res.setContent(hContent);
                }
            }
            resList.add(res);
        }
        return resList;
    }

    /**
     * 分词搜索资源信息，无高亮
     * @param q 查询关键字
     * @return
     * @throws Exception
     */
    public List<Res> searchNoHighLighter(String q) throws Exception {
        dir = FSDirectory.open(Paths.get(lucenePath));
        IndexReader reader = DirectoryReader.open(dir);
        IndexSearcher is = new IndexSearcher(reader);

        BooleanQuery.Builder booleanQuery = new BooleanQuery.Builder();
        SmartChineseAnalyzer analyzer = new SmartChineseAnalyzer();
        QueryParser parserName = new QueryParser("name", analyzer);
        Query queryName = parserName.parse(q);
        booleanQuery.add(queryName, BooleanClause.Occur.SHOULD);

        QueryParser parserContent = new QueryParser("content", analyzer);
        Query queryContent = parserContent.parse(q);
        booleanQuery.add(queryContent, BooleanClause.Occur.SHOULD);

        TopDocs hits = is.search(booleanQuery.build(), 20);
        List<Res> resList = new LinkedList<>();
        for (ScoreDoc scoreDoc : hits.scoreDocs) {
            Document doc = is.doc(scoreDoc.doc);
            Res res = new Res();
            res.setId(Integer.parseInt(doc.get(("id"))));
            String name = doc.get("name");
            res.setName(name);
            resList.add(res);
        }
        return resList;
    }
}
